from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain import hub
from langchain_core.tools import tool
from langchain_ollama import ChatOllama

from common_config import CHAT_OLLAMA_MODEL


# Step 1: 定义一个简单的工具（比如获取天气信息）
@tool
def get_weather(city: str) -> str:
    """返回指定城市的当前天气情况（模拟数据）"""
    weather_data = {
        "北京": "晴天，25°C",
        "上海": "多云，28°C",
        "广州": "雷阵雨，30°C"
    }
    return weather_data.get(city, "未知城市或无数据")


# Step 2: 定义另一个工具（比如计算器）
@tool
def calculator(expression: str) -> str:
    """计算数学表达式（简单示例）"""
    try:
        result = eval(expression)
        return f"结果是：{result}"
    except Exception as e:
        return f"错误：{e}"


# Step 3: 初始化 Ollama 模型
model = ChatOllama(model=CHAT_OLLAMA_MODEL, temperature=0)


# Step 4: 加载预定义的 prompt（包含工具调用格式说明）
prompt = hub.pull("hwchase17/openai-functions-agent")


# Step 5: 创建 agent 和绑定工具
tools = [get_weather, calculator]
agent = create_tool_calling_agent(model, tools, prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)


# Step 6: 定义一个函数来处理用户输入并执行 agent
def run_agent(query):
    response = agent_executor.invoke({
        "input": query
    })
    print("回答：", response["output"])


# Step 7: 示例运行
if __name__ == "__main__":
    print("欢迎使用 LangChain + Ollama Tool Call 示例")
    while True:
        user_input = input("\n请输入问题（输入 'exit' 退出）：")
        if user_input.lower() in ["exit", "quit"]:
            break
        run_agent(user_input)